package com.xyhy.core.utils.jdk;

import cn.hutool.core.lang.Opt;
import cn.hutool.core.stream.CollectorUtil;
import cn.hutool.core.stream.SimpleCollector;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;

public class CollectorUtils extends CollectorUtil {
  public static <T, K, U> Collector<T, ?, Map<K, U>> toMap(Function<? super T, ? extends K> keyMapper, //
                                                           Function<? super T, ? extends U> valueMapper) {
    return toMap(keyMapper, valueMapper, (x1, x2) -> Optional.ofNullable(x2).orElse(x1));
  }

  public static <T, K, U> Collector<T, ?, Map<K, U>> toMap(Function<? super T, ? extends K> keyMapper,
                                                           Function<? super T, ? extends U> valueMapper,
                                                           BinaryOperator<U> mergeFunction) {
    return toMap(keyMapper, valueMapper, mergeFunction, HashMap::new);
  }

  /**
   * 对null友好的 toMap 操作的 {@link Collector}实现
   *
   * @param keyMapper     指定map中的key
   * @param valueMapper   指定map中的value
   * @param mergeFunction 合并前对value进行的操作
   * @param mapSupplier   最终需要的map类型
   * @param <T>           实体类型
   * @param <K>           map中key的类型
   * @param <U>           map中value的类型
   * @param <M>           map的类型
   * @return 对null友好的 toMap 操作的 {@link Collector}实现
   */
  public static <T, K, U, M extends Map<K, U>>
  Collector<T, ?, M> toMap(Function<? super T, ? extends K> keyMapper,
                           Function<? super T, ? extends U> valueMapper,
                           BinaryOperator<U> mergeFunction,
                           Supplier<M> mapSupplier) {
    BiConsumer<M, T> accumulator
            = (map, element) -> map.put(Opt.ofNullable(element).map(keyMapper).get(),
            Opt.ofNullable(element).map(valueMapper).get());
    return new SimpleCollector<>(mapSupplier, accumulator, mapMerger(mergeFunction), CH_ID);
  }
}
